#include "Spectrum.h"

using namespace std;
using namespace clUtil;

static const float kMassH = 1.008f;

void Spectrum::processRun(xmlNode* msRunNode)
{
  xmlNode* curNode;
  xmlChar* curNodePropertyValue;
  mMaxScanSize = 0;
  mCurrentScan = 0;

  //Get the total number of scans
  curNodePropertyValue = xmlGetProp(msRunNode, (const xmlChar*)"scanCount");
  sscanf((const char*)curNodePropertyValue, "%d", &mTotalScans);

  for(curNode = msRunNode->children; curNode != NULL; curNode = curNode->next)
  {
    if(strcmp((const char*)curNode->name, "scan") == 0)
    {
      //Allocate a new scan
      unique_ptr<Scan> newScan(new Scan(curNode, mMaxScanSize));

      //Give ownership to the scan vector
      mMSScans.push_back(move(newScan));
    }
  }

  //Allocate all our flat buffers
  unique_ptr<unsigned int[]> 
    allPeaks(new unsigned int[mMaxScanSize * mTotalScans]);
  unique_ptr<unsigned int[]> 
    allMZ(new unsigned int[mMaxScanSize * mTotalScans]);
  unique_ptr<MSScanInfo[]> allScanInfo(new MSScanInfo[mTotalScans]);

  mAllPeaks = move(allPeaks);
  mAllMZ = move(allMZ);
  mAllScanInfo = move(allScanInfo);

  unsigned int* curPeaksPtr = mAllPeaks.get();
  unsigned int* curMZPtr = mAllMZ.get();
  MSScanInfo* curScanInfoPtr = mAllScanInfo.get();

  //Create a flat array of all scans
  for(size_t i = 0; i < mMSScans.size(); i++)
  {
    mMSScans[i]->fillFlatArrays(curPeaksPtr,
                                curMZPtr,
                                curScanInfoPtr,
                                mMaxScanSize);
  }
}

Spectrum::Spectrum(const char* fileName) : mSpecName(fileName)
{
  xmlDoc* mzxmlFile;
  xmlNode* xmlRoot;
  xmlNode* curNode;
 
  mBaseName = mSpecName;
  mBaseName.erase(mBaseName.find(".mzXML"));

  mzxmlFile = xmlReadFile(fileName, NULL, 0);

  if(mzxmlFile == NULL)
  {
    throw SpectrumException("Could not open mzxmlfile");
  }

  xmlRoot = xmlDocGetRootElement(mzxmlFile);

  for(curNode = xmlRoot->children; curNode != NULL; curNode = curNode->next)
  {
    if(strcmp((const char*)curNode->name, "msRun") == 0)
    {
      this->processRun(curNode);
      break;
    }
  }

  xmlFreeDoc(mzxmlFile);

  //Preallocate space for scores
  mCombinedScores.resize(mTotalScans);

  for(size_t curScan = 0; curScan < mTotalScans; curScan++)
  {
    mCombinedScores[curScan].resize(gMaxPrecursorCharge);
  }
}

void Spectrum::allocateSpectrum(size_t currentDevice)
{
  if(gDeviceAlgorithmConfig[currentDevice].SpectrumStorage != 
     kStorageMechanismImage)
  {
    //Allocate mz buffer
    unique_ptr<Memory> mzBuffer(new Buffer(kScansPerIteration * 
                                           sizeof(unsigned int) *
                                           gMaxPeakCount));

    mDeviceFilteredMZ.push_back(move(mzBuffer));
    
    //Allocate peaks buffer
    unique_ptr<Memory> peaksBuffer(new Buffer(kScansPerIteration * 
                                              sizeof(unsigned int) *
                                              gMaxPeakCount));

    mDeviceFilteredPeaks.push_back(move(peaksBuffer));
  }
  else
  {
    cl_image_format format;
    size_t imageWidth = gMaxPeakCount; 

    format.image_channel_data_type = CL_UNSIGNED_INT32;
    format.image_channel_order = CL_RGBA; 

    //Allocate mz image
    unique_ptr<Memory> mzImage(new Image(imageWidth,
                                         kScansPerIteration,
                                         format.image_channel_order,
                                         format.image_channel_data_type));
    
    mDeviceFilteredMZ.push_back(move(mzImage));
    
    //Allocate peaks image
    unique_ptr<Memory> peaksImage(new Image(imageWidth,
                                            kScansPerIteration,
                                            format.image_channel_order,
                                            format.image_channel_data_type));
    
    mDeviceFilteredPeaks.push_back(move(peaksImage));
  }
}

void Spectrum::allocateDevices()
{
  //Allocate a ton of fucking shit on each device
  for(unsigned int curDevice = 0; 
      curDevice < Device::GetDevices().size(); 
      curDevice++)
  {
    Device::SetCurrentDevice(curDevice);

    //Allocate spectrum data
    allocateSpectrum(curDevice);

    size_t byteCount = sizeof(unsigned int) * mMaxScanSize * mTotalScans;

    unique_ptr<Memory> tmpBuffer;

    //Allocate more shit
    tmpBuffer = move(unique_ptr<Memory>(new Buffer(byteCount)));
    mDevicePeaks.push_back(move(tmpBuffer));
    
    tmpBuffer = move(unique_ptr<Memory>(new Buffer(byteCount)));
    mDevicePeaksTmp.push_back(move(tmpBuffer));
    
    tmpBuffer = move(unique_ptr<Memory>(new Buffer(byteCount)));
    mDeviceMZ.push_back(move(tmpBuffer));
    
    tmpBuffer = move(unique_ptr<Memory>(new Buffer(byteCount)));
    mDeviceMZTmp.push_back(move(tmpBuffer));
    
    tmpBuffer = move(unique_ptr<Memory>(new Buffer(sizeof(MSScanInfo) * 
                                                   kScansPerIteration)));
    mDeviceScanInfo.push_back(move(tmpBuffer));
 
    tmpBuffer = move(unique_ptr<Memory>(new Buffer(sizeof(cl_float) * 
                                                   kScansPerIteration * 
                                                   gBestScoreListSize *
                                                   gMaxPrecursorCharge)));
    mDeviceScores.push_back(move(tmpBuffer));

    tmpBuffer = move(unique_ptr<Memory>(new Buffer(sizeof(cl_uint) * 
                                                   kScansPerIteration * 
                                                   gBestScoreListSize *
                                                   gMaxPrecursorCharge)));
    mDeviceScoresIdx.push_back(move(tmpBuffer));
  }
}

void Spectrum::freeDevices()
{
  mDeviceFilteredMZ.clear();
  mDeviceFilteredPeaks.clear();
  mDevicePeaks.clear();
  mDevicePeaksTmp.clear();
  mDeviceMZ.clear();
  mDeviceMZTmp.clear();
  mDeviceScanInfo.clear();
  mDeviceScores.clear();
  mDeviceScoresIdx.clear();
}

void Spectrum::processSpectrum(Database& peptides)
{
  size_t byteCount;

  vector<float> 
    scoresHost(mTotalScans * gBestScoreListSize * gMaxPrecursorCharge, 0.0f);
  vector<unsigned int> 
    scoresIdxHost(mTotalScans * gBestScoreListSize * gMaxPrecursorCharge, 0u);

  double time1 = Utility::getTime();

  //Score scans in parallel using all OpenCL devices
  ParallelFor(0, 1, mTotalScans - 1, [&](size_t start, size_t end)
  {
    size_t deviceID = Device::GetCurrentDeviceNum();

    //cout << start << " " << end << endl;

    unsigned int numIterations = end - start + 1;
    size_t maxScanSize = mMaxScanSize * sizeof(unsigned int);

    //Break index range into chunks that fit in memory
    for(unsigned int curScan = 0; 
        curScan < numIterations; 
        curScan += kScansPerIteration)
    {
      unsigned int scansToDo = numIterations - curScan > kScansPerIteration ?
        kScansPerIteration :
        numIterations - curScan;

      byteCount = maxScanSize * scansToDo;

      //Get spectrum device pointers
      Memory* peaks = mDevicePeaks[deviceID].get();
      Memory* mz = mDeviceMZ[deviceID].get();
      Memory* peaksTmp = mDevicePeaksTmp[deviceID].get();
      Memory* mzTmp = mDeviceMZTmp[deviceID].get();
      Memory* scanInfo = mDeviceScanInfo[deviceID].get();
      Memory* filteredMZ = mDeviceFilteredMZ[deviceID].get();
      Memory* filteredPeaks = mDeviceFilteredPeaks[deviceID].get();
      Memory* scores = mDeviceScores[deviceID].get();
      Memory* scoresIdx = mDeviceScoresIdx[deviceID].get();

      //Get database device pointers
      Memory* peptideStrings = peptides.getPeptides();
      Memory* peptideMasses = peptides.getMasses();

      unsigned int maxScanSizeUI = mMaxScanSize;
      unsigned int numPeptides = peptides.getNumPeptides();

      //Transfer peaks, mz, and scan headers to devicea
      peaks->put(&mAllPeaks[(start + curScan) * mMaxScanSize], byteCount);
      mz->put(&mAllMZ[(start + curScan) * mMaxScanSize], byteCount);
      scanInfo->put(&mAllScanInfo[start + curScan], 
                    sizeof(MSScanInfo) * scansToDo);
      
      size_t blockSize = gDeviceAlgorithmConfig[deviceID].LocalWorkSize;

      //Preprocess
      clUtilEnqueueKernel("preprocessSpectra",
                          clUtilGrid(blockSize, blockSize, scansToDo, 1),
                          peaks,
                          mz,
                          peaksTmp,
                          mzTmp,
                          maxScanSizeUI,
                          scanInfo);

      //Pack spectra into image
      clUtilEnqueueKernel("pack",
                          clUtilGrid(blockSize, blockSize, scansToDo, 1),
                          mz,
                          peaks,
                          maxScanSizeUI,
                          scanInfo,
                          filteredMZ,
                          filteredPeaks);

      //Process the spectra
      clUtilEnqueueKernel("matchSpectra",
                          clUtilGrid(blockSize, blockSize, scansToDo, 1),
                          filteredMZ,
                          filteredPeaks,
                          scanInfo,
                          peptideStrings,
                          peptideMasses,
                          numPeptides,
                          scores,
                          scoresIdx);

      //Fetch the scores
      scores->get(&scoresHost[start * 
                              gBestScoreListSize * 
                              gMaxPrecursorCharge], 
                  sizeof(cl_float) * 
                  scansToDo * gBestScoreListSize * gMaxPrecursorCharge);
      scoresIdx->get(&scoresIdxHost[start * 
                                    gBestScoreListSize * 
                                    gMaxPrecursorCharge],
                     sizeof(cl_uint) * 
                     scansToDo * gBestScoreListSize * gMaxPrecursorCharge);
    }
  });

  for(unsigned int curScan = 0; curScan < mTotalScans; curScan++)
  {
    for(unsigned int curCharge = 0; 
        curCharge < gMaxPrecursorCharge; 
        curCharge++)
    {
      for(unsigned int curScore = 0; curScore < gBestScoreListSize; curScore++)
      {
        ScorePair temp;
        unsigned int index = curScan *
                             gBestScoreListSize *
                             gMaxPrecursorCharge +
                             curCharge *
                             gBestScoreListSize +
                             curScore;
        
        temp.score = scoresHost[index];
        temp.idx = scoresIdxHost[index];

        mCombinedScores[curScan][curCharge].push_back(temp);
      }
    }
  }

  double time2 = Utility::getTime();

  cout << "Time:\t" << time2 - time1 << endl; 
}

void Spectrum::writeHeader(FILE* outputFile)
{
  fprintf(outputFile, "<?xml version=\"1.0\" encoding=\"ISO-8859-1\"?>\n");

  time_t curTime = time(0);
  struct tm* localTime = localtime(&curTime);
  string asciiTime(asctime(localTime));

  //Remove trailing \n from timestamp
  asciiTime[asciiTime.length() - 1] = '\0';

  fprintf(outputFile, 
          "<msms_pipeline_analysis date=\"%s\" summary_xml=\"%s\" "
          "xmlns=\"http://regis-web.systemsbiology.net/pepXML\" "
          "xmlns:xsi=\"http://www.w3.org/2001/XMLSchema-instance\" "
          "xsi:schemaLocation=\"http://sashimi.sourceforge.net/schema_revision/"
          "pepXML/pepXML_v117.xsd\">\n",
          asctime(localTime),
          string(mBaseName + ".pepXML").c_str());

  fprintf(outputFile,
          "\t<analysis_summary analysis=\"Specmaster\" version=\"%s\" "
          "time=\"%s\"/>\n",
          gVersion,
          asciiTime.c_str());

  fprintf(outputFile,
          "\t<msms_run_summary base_name=\"%s\" raw_data_type=\"\" "
          "raw_data=\"\">\n",
          mBaseName.c_str());

  fprintf(outputFile, 
          "\t\t<sample_enzyme name=\"Trypsin/P\" "
          "independent=\"false\" fidelity=\"specific\">\n");

  fprintf(outputFile, 
          "\t\t\t<specificity sense=\"C\" cut=\"KR\" no_cut=\"\" "
          "min_spacing=\"1\"/>\n");

  fprintf(outputFile, "\t\t</sample_enzyme>\n");

  fprintf(outputFile, 
          "\t\t<search_summary base_name=\"%s\" search_engine=\"Specmaster\" "
          "precursor_mass_type=\"monoisotopic\" "
          "fragment_mass_type=\"monoisotopic\" "
          "out_data_type=\"\" out_data=\"\">\n",
          mBaseName.c_str());
  fprintf(outputFile, "\t\t</search_summary>\n");

}

void Spectrum::combineScores(Database& peptides, const char* outputFileName)
{
  string outputName = mSpecName.replace(mSpecName.rfind("."), 
                                        7, 
                                        string(".pepXML"));

  FILE* outputFile = fopen(outputName.c_str(), "w");
  unsigned int curQuery = 0;

  if(outputFile == NULL)
  {
    printf("ERROR: Could not open output file for writing.\n");
    exit(1);
  }

  //Write results to file
  this->writeHeader(outputFile);

  for(unsigned int curScan = 0; curScan < mTotalScans; curScan++)
  {
    for(unsigned int curCharge = 0; 
        curCharge < gMaxPrecursorCharge; 
        curCharge++)
    {

      //sort by score 
      sort(mCombinedScores[curScan][curCharge].begin(), 
           mCombinedScores[curScan][curCharge].end());

      //If we have scores for that charge state, write a query
      if(mCombinedScores[curScan][curCharge][0].score != 0.0f)
      {
        fprintf(outputFile, 
                "\t\t<search_query spectrum=\"%s.%u.%u.%u\" "
                "spectrumNativeID=\"controllerType=0 controllerNumber=1 "
                "scan=%u\" start_scan=\"%u\" end_scan=\"%u\" "
                "precursor_neutral_mass=\"%f\" assumed_charge=\"%u\""
                "index=\"%u\">\n",
                mBaseName.c_str(),
                curScan + 1,
                curScan + 1,
                curCharge,
                curScan + 1,
                curScan + 1,
                curScan + 1,
                (float)mAllScanInfo[curScan].mPrecursorMZ * 
                  (float)curCharge / 100.0f - (float)curCharge * kMassH,
                curCharge + 1,
                curQuery + 1);
            
        fprintf(outputFile, "\t\t\t<search_result>\n");

        curQuery++;
      }
      else
      {
        continue;
      }

      if(mAllScanInfo[curScan].mMSLevel != 1)
      {
        for(unsigned int curScore = 0; 
            curScore < gBestScoreListSize; 
            curScore++)
        {
          if(mCombinedScores[curScan][curCharge][curScore].score != 0.0f)
          {
            ScorePair& curScoreRef = 
              mCombinedScores[curScan][curCharge][curScore];
            
            string& peptideSequence = 
              peptides.getPeptide(curScoreRef.idx).mAcids;

            pair<multimap<string, string>::iterator,
              multimap<string, string>::iterator> proteins = 
                peptides.getProtein(peptideSequence);


            fprintf(outputFile,
                    "\t\t\t\t<search_hit hit_rank=\"%u\" peptide=\"%s\" "
                    "protein=\"%s\">\n",
                    curScore + 1,
                    peptideSequence.c_str(),
                    proteins.first->second.c_str());

            auto curProtein = proteins.first;
            curProtein++;

            for( ; curProtein != proteins.second; curProtein++)
            {
              fprintf(outputFile,
                      "\t\t\t\t\t<alternative_protein protein\"%s\"/>\n",
                      curProtein->second.c_str());
            }
            fprintf(outputFile,
                    "\t\t\t\t\t<search_score name=\"mvh\" value=\"%f\"/>\n",
                    mCombinedScores[curScan][curCharge][curScore].score);

            fprintf(outputFile, "\t\t\t\t</search_hit>\n");
          }
        }
      }

      fprintf(outputFile, "\t\t\t</search_result>\n");
      fprintf(outputFile, "\t\t</search_query>\n");
    }
  }
    fprintf(outputFile, "\t</msms_run_summary>\n");
    fprintf(outputFile, "</msms_pipeline_analysis>\n");

    fclose(outputFile);
  }


void Spectrum::printScan(size_t id)
{
  printf("###SCAN NUM %lu\n", id);

  for(unsigned int i = 0; i < mAllScanInfo[id].mNumPeaks; i++)
  {
    printf("%u\t%u\n",
           mAllMZ[id * mMaxScanSize + i],
           mAllPeaks[id * mMaxScanSize + i]);
  }
}

void Spectrum::printAllScans()
{
  for(unsigned int i = 0 ; i < mTotalScans; i++)
  {
    this->printScan(i);
  }

  exit(0);
}
